import pandas as pd
import os
from typing import Union, Dict, Any, List, Optional
class FileDataEngine:
    def __init__(
        self,
        file_path: str,
        table_name: Optional[str] = None,
        return_format: str = "json",
        auto_create_table: bool = True,
        auto_add_id: bool = True,
    ):
        """
        初始化 FileDataEngine。
        :param file_path: 文件路径。
        :param table_name: 默认表名。
        :param return_format: 返回格式（json 或 DataFrame）。
        :param auto_create_table: 是否自动创建表。
        :param auto_add_id: 是否自动添加 ID 列。
        """
        self.file_path = file_path
        self.return_format = return_format
        self.auto_create_table = auto_create_table
        self.auto_add_id = auto_add_id
        self.id_column = "id"  # 默认主键列名

        # 根据文件类型生成默认表名
        _, ext = os.path.splitext(file_path)
        if ext.lower() == ".csv":
            # CSV 文件不支持传入 table_name，默认使用文件名
            self.default_table_name = os.path.splitext(os.path.basename(file_path))[0]
        else:
            # 其他文件类型（如 Excel）支持传入 table_name
            self.default_table_name = table_name if table_name is not None else os.path.splitext(os.path.basename(file_path))[0]

        # 初始化 tables
        if not os.path.exists(file_path):
            if auto_create_table:
                columns = [self.id_column] if self.auto_add_id else []
                self.tables = {self.default_table_name: pd.DataFrame(columns=columns)}
            else:
                raise ValueError(f"文件 '{file_path}' 不存在，且自动创建文件已关闭。")
        else:
            self.tables = self._read_from_file(file_path)

    def _read_from_file(self, file_path: str) -> Dict[str, pd.DataFrame]:
        """
        从文件中读取数据。
        """
        _, ext = os.path.splitext(file_path)
        ext = ext.lower()
        if ext == ".csv":
            # CSV 文件只支持单表，忽略传入的 table_name
            df = pd.read_csv(file_path)
            if self.auto_add_id and self.id_column not in df.columns:
                df[self.id_column] = range(1, len(df) + 1)
            return {self.default_table_name: df}
        elif ext in (".xls", ".xlsx"):
            # Excel 文件支持多表
            excel_file = pd.ExcelFile(file_path)
            sheet_names = excel_file.sheet_names
            tables = {}
            for sheet_name in sheet_names:
                df = excel_file.parse(sheet_name)
                if self.auto_add_id and self.id_column not in df.columns:
                    df[self.id_column] = range(1, len(df) + 1)
                tables[sheet_name] = df
            return tables
        elif ext == ".json":
            # JSON 文件只支持单表
            df = pd.read_json(file_path)
            if self.auto_add_id and self.id_column not in df.columns:
                df[self.id_column] = range(1, len(df) + 1)
            return {self.default_table_name: df}
        else:
            raise ValueError(f"不支持的文件类型: {ext}")

    def _persist_to_file(self):
        """
        将数据持久化到文件。
        """
        _, ext = os.path.splitext(self.file_path)
        ext = ext.lower()
        if ext == ".csv":
            # CSV 文件只保存默认表
            df = self.tables.get(self.default_table_name, pd.DataFrame())
            df.to_csv(self.file_path, index=False)
        elif ext in (".xls", ".xlsx"):
            # Excel 文件保存所有表
            with pd.ExcelWriter(self.file_path) as writer:
                for sheet_name, df in self.tables.items():
                    df.to_excel(writer, sheet_name=sheet_name, index=False)
        elif ext == ".json":
            # JSON 文件只保存默认表
            df = self.tables.get(self.default_table_name, pd.DataFrame())
            df.to_json(self.file_path, orient='records', lines=True)
        else:
            raise ValueError(f"不支持的文件类型: {ext}")

    def save(self, entity: Dict[str, Any], table_name: Optional[str] = None) -> str:
        """
        插入数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        # 对于 CSV 文件，忽略传入的 table_name，始终使用默认表名
        _, ext = os.path.splitext(self.file_path)
        if ext.lower() == ".csv":
            table_name = self.default_table_name

        # 如果表不存在，且允许自动创建表
        if table_name not in self.tables:
            if self.auto_create_table:
                columns = [self.id_column] if self.auto_add_id else []
                self.tables[table_name] = pd.DataFrame(columns=columns)
            else:
                raise ValueError(f"表 '{table_name}' 不存在，且自动建表已关闭。")

        # 追加数据
        df = self.tables[table_name]
        new_row = pd.DataFrame([entity])
        df = pd.concat([df, new_row], ignore_index=True)

        # 自动生成 ID
        if self.auto_add_id and self.id_column in df.columns:
            if df[self.id_column].isnull().any():
                # 计算当前最大 ID，如果 DataFrame 为空，则 max_id 为 0
                max_id = df[self.id_column].max(skipna=True)
                if pd.isna(max_id):  # 如果 max_id 是 NaN，则设置为 0
                    max_id = 0
                # 为缺失 ID 的行生成新 ID
                df.loc[df[self.id_column].isnull(), self.id_column] = range(int(max_id) + 1, int(max_id) + 1 + df[
                    self.id_column].isnull().sum())

        self.tables[table_name] = df  # 更新表数据
        self._persist_to_file()  # 持久化到文件
        return str(df.iloc[-1][self.id_column])  # 返回插入的 ID

    def delete(self, query: Dict[str, Any], table_name: Optional[str] = None) -> int:
        """
        删除数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]
        # 根据查询条件删除数据
        for key, value in query.items():
            df = df[df[key] != value]
        deleted_count = len(self.tables[table_name]) - len(df)
        self.tables[table_name] = df  # 更新表数据
        self._persist_to_file()  # 持久化到文件
        return deleted_count

    def deleteById(self, id: str, table_name: Optional[str] = None) -> int:
        """
        根据 ID 删除数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]
        # 根据 ID 删除数据
        df = df[df[self.id_column] != id]
        deleted_count = len(self.tables[table_name]) - len(df)
        self.tables[table_name] = df  # 更新表数据
        self._persist_to_file()  # 持久化到文件
        return deleted_count

    def update(self, query: Dict[str, Any], update_data: Dict[str, Any], table_name: Optional[str] = None) -> int:
        """
        更新数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]
        # 根据查询条件更新数据
        for key, value in query.items():
            df.loc[df[key] == value, list(update_data.keys())] = list(update_data.values())
        self.tables[table_name] = df  # 更新表数据
        self._persist_to_file()  # 持久化到文件
        return len(df[df[key] == value])  # 返回更新的行数

    def updateById(self, id: str, update_data: Dict[str, Any], table_name: Optional[str] = None) -> int:
        """
        根据 ID 更新数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]
        # 根据 ID 更新数据
        df.loc[df[self.id_column] == id, list(update_data.keys())] = list(update_data.values())
        self.tables[table_name] = df  # 更新表数据
        self._persist_to_file()  # 持久化到文件
        return 1  # 返回更新的行数（ID 唯一，所以总是 1）

    def find(self, query: Dict[str, Any], table_name: Optional[str] = None) -> Union[List[Dict], pd.DataFrame]:
        """
        查询数据。
        :param query: 查询条件字典。
        :param table_name: 表名。
        :return: 查询结果（列表或 DataFrame）。
        :raises ValueError: 如果表不存在或字段不存在。
        """
        if table_name is None:
            table_name = self.default_table_name

        # 检查表是否存在
        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]

        # 检查查询条件中的字段是否存在
        for key in query.keys():
            if key not in df.columns:
                raise ValueError(f"表 '{table_name}' 中没有 '{key}' 列。")

        # 遍历查询条件，过滤数据
        for key, value in query.items():
            df = df[df[key] == value]

        if self.return_format == "json":
            return df.to_dict(orient="records")
        else:
            return df

    def findAll(self, table_name: Optional[str] = None) -> Union[List[Dict], pd.DataFrame]:
        """
        查询集合中的所有数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]
        if self.return_format == "json":
            return df.to_dict(orient="records")
        else:
            return df

    def findById(self, id: str, table_name: Optional[str] = None) -> Union[Dict, pd.DataFrame]:
        """
        根据 ID 查询数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]
        result = df[df[self.id_column] == id]
        if result.empty:
            return None

        if self.return_format == "json":
            return result.iloc[0].to_dict()
        else:
            return result

    def findPage(self, query: Dict[str, Any], page: int, page_size: int, table_name: Optional[str] = None) -> Union[List[Dict], pd.DataFrame]:
        """
        分页查询数据。
        """
        if table_name is None:
            table_name = self.default_table_name

        if table_name not in self.tables:
            raise ValueError(f"表 '{table_name}' 不存在。")

        df = self.tables[table_name]
        # 根据查询条件过滤数据
        for key, value in query.items():
            df = df[df[key] == value]

        # 分页
        skip = (page - 1) * page_size
        df = df.iloc[skip:skip + page_size]

        if self.return_format == "json":
            return df.to_dict(orient="records")
        else:
            return df